# Copyright      2023  Xiaomi Corp.        (authors: Zengrui Jin)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import logging
import os
from functools import lru_cache
from pathlib import Path
from typing import Dict

from lhotse import CutSet, load_manifest_lazy


class MultiDataset:
    def __init__(self, args: argparse.Namespace):
        """
        Args:
          manifest_dir:
            It is expected to contain the following files:
            - aishell2_cuts_train.jsonl.gz
        """

        # self.fbank_dir  = Path("/home/work_nfs8/xlgeng/new_workspace/icefall/egs/multi_zh_en/ASR/gxl_data/3000h/fbank_common")
        self.train_fbank_dir = Path(args.train_fbank_dir)
        self.valid_fbank_dir = Path(args.valid_fbank_dir)
        self.test_parent_dir = Path(args.test_parent_dir)
        self.train_prefix = args.train_prefix
        self.valid_prefix = args.valid_prefix

    def train_cuts(self) -> CutSet:
        logging.info("耿雪龙: About to get multidataset train cuts")
        all_data_cuts = load_manifest_lazy(
            self.train_fbank_dir / f"{self.train_prefix}_cuts_train.jsonl.gz"
        )
        # wenetspeech_half_cuts =  load_manifest_lazy(
        #     self.train_fbank_dir / f"{self.train_prefix}_cuts_train_half.jsonl.gz"
        # )
        return all_data_cuts


    def dev_cuts(self) -> CutSet:
        logging.info("About to get multidataset dev cuts")
        gxl_dev_cuts = load_manifest_lazy(
            self.valid_fbank_dir / f"{self.valid_prefix}_cuts_dev.jsonl.gz"
        )
        return gxl_dev_cuts


    def test_cuts(self) -> Dict[str, CutSet]:
        logging.info("About to get multidataset test cuts")
        test_family_dict = {}
        test_parent_dir_str = str(self.test_parent_dir)
        dataname_list = os.listdir(test_parent_dir_str)
        for dataname in dataname_list:
            prefix_i = dataname
            fbank_dir_i = os.path.join(test_parent_dir_str, dataname, "fbank")
            cut_path_i = Path(fbank_dir_i) / f"{prefix_i}_cuts_test.jsonl.gz"
            if not os.path.exists(cut_path_i):
                continue
            cuts_i = load_manifest_lazy(
                cut_path_i
            )
            test_family_dict[dataname] = cuts_i
        return test_family_dict